package org.gwh.airagknowledge.repository;

import org.gwh.airagknowledge.entity.VectorEmbedding;
import org.springframework.context.annotation.Primary;
import org.springframework.data.domain.Example;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;
import org.springframework.data.repository.query.FluentQuery;
import org.springframework.stereotype.Repository;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 测试专用的VectorEmbedding仓库模拟实现
 * 使用内存数据结构代替实际数据库，避免PGvector依赖
 */
@Repository
@Primary
public class MockVectorEmbeddingRepository implements VectorEmbeddingRepository {

    private final Map<Long, VectorEmbedding> embeddingStore = new ConcurrentHashMap<>();
    private final AtomicLong idSequence = new AtomicLong(1);

    @Override
    public List<VectorEmbedding> findAll() {
        return new ArrayList<>(embeddingStore.values());
    }

    @Override
    public List<VectorEmbedding> findAll(Sort sort) {
        return findAll();
    }

    @Override
    public Page<VectorEmbedding> findAll(Pageable pageable) {
        throw new UnsupportedOperationException("Not implemented for mock");
    }

    @Override
    public List<VectorEmbedding> findAllById(Iterable<Long> ids) {
        List<VectorEmbedding> result = new ArrayList<>();
        ids.forEach(id -> {
            if (embeddingStore.containsKey(id)) {
                result.add(embeddingStore.get(id));
            }
        });
        return result;
    }

    @Override
    public long count() {
        return embeddingStore.size();
    }

    @Override
    public void deleteById(Long id) {
        embeddingStore.remove(id);
    }

    @Override
    public void delete(VectorEmbedding entity) {
        if (entity.getId() != null) {
            embeddingStore.remove(entity.getId());
        }
    }

    @Override
    public void deleteAllById(Iterable<? extends Long> ids) {
        ids.forEach(embeddingStore::remove);
    }

    @Override
    public void deleteAll(Iterable<? extends VectorEmbedding> entities) {
        entities.forEach(entity -> {
            if (entity.getId() != null) {
                embeddingStore.remove(entity.getId());
            }
        });
    }

    @Override
    public void deleteAll() {
        embeddingStore.clear();
    }

    @Override
    public <S extends VectorEmbedding> S save(S entity) {
        if (entity.getId() == null) {
            entity.setId(idSequence.getAndIncrement());
        }
        embeddingStore.put(entity.getId(), entity);
        return entity;
    }

    @Override
    public <S extends VectorEmbedding> List<S> saveAll(Iterable<S> entities) {
        List<S> savedEntities = new ArrayList<>();
        entities.forEach(entity -> savedEntities.add(save(entity)));
        return savedEntities;
    }

    @Override
    public Optional<VectorEmbedding> findById(Long id) {
        return Optional.ofNullable(embeddingStore.get(id));
    }

    @Override
    public boolean existsById(Long id) {
        return embeddingStore.containsKey(id);
    }

    @Override
    public void flush() {
        // No-op for mock implementation
    }

    @Override
    public <S extends VectorEmbedding> S saveAndFlush(S entity) {
        return save(entity);
    }

    @Override
    public <S extends VectorEmbedding> List<S> saveAllAndFlush(Iterable<S> entities) {
        return saveAll(entities);
    }

    @Override
    public void deleteAllInBatch(Iterable<VectorEmbedding> entities) {
        deleteAll(entities);
    }

    @Override
    public void deleteAllByIdInBatch(Iterable<Long> ids) {
        deleteAllById(ids);
    }

    @Override
    public void deleteAllInBatch() {
        deleteAll();
    }

    @Override
    public VectorEmbedding getOne(Long id) {
        return findById(id).orElse(null);
    }

    @Override
    public VectorEmbedding getById(Long id) {
        return findById(id).orElse(null);
    }

    @Override
    public VectorEmbedding getReferenceById(Long id) {
        return findById(id).orElse(null);
    }

    @Override
    public <S extends VectorEmbedding> Optional<S> findOne(Example<S> example) {
        throw new UnsupportedOperationException("Not implemented for mock");
    }

    @Override
    public <S extends VectorEmbedding> List<S> findAll(Example<S> example) {
        throw new UnsupportedOperationException("Not implemented for mock");
    }

    @Override
    public <S extends VectorEmbedding> List<S> findAll(Example<S> example, Sort sort) {
        throw new UnsupportedOperationException("Not implemented for mock");
    }

    @Override
    public <S extends VectorEmbedding> Page<S> findAll(Example<S> example, Pageable pageable) {
        throw new UnsupportedOperationException("Not implemented for mock");
    }

    @Override
    public <S extends VectorEmbedding> long count(Example<S> example) {
        throw new UnsupportedOperationException("Not implemented for mock");
    }

    @Override
    public <S extends VectorEmbedding> boolean exists(Example<S> example) {
        throw new UnsupportedOperationException("Not implemented for mock");
    }

    @Override
    public <S extends VectorEmbedding, R> R findBy(Example<S> example, Function<FluentQuery.FetchableFluentQuery<S>, R> queryFunction) {
        throw new UnsupportedOperationException("Not implemented for mock");
    }

    @Override
    public List<VectorEmbedding> findSimilarEmbeddings(Long knowledgeBaseId, Object queryEmbedding, int limit) {
        // 简单实现，不考虑向量相似度，只返回与知识库相关的前limit个结果
        return embeddingStore.values().stream()
                .limit(limit)
                .collect(Collectors.toList());
    }

    @Override
    public void deleteByChunkId(Long chunkId) {
        // 找到所有与chunkId相关的嵌入并删除它们
        embeddingStore.values().stream()
                .filter(embedding -> embedding.getChunk() != null && 
                        embedding.getChunk().getId().equals(chunkId))
                .map(VectorEmbedding::getId)
                .forEach(embeddingStore::remove);
    }
} 